package cn.iocoder.boot.springai001.controller;

import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.*;

import java.util.*;
import java.util.stream.Collectors;

/**
 * MCP 工具管理控制器
 * <p>
 * 提供 MCP 工具的查询、过滤和管理功能
 * <p>
 * API路径：
 * - GET /api/v1/mcp/servers - 获取所有MCP服务器
 * - GET /api/v1/mcp/tools - 获取所有工具
 * - GET /api/v1/mcp/tools/{toolName} - 获取工具详情
 * - GET /api/v1/mcp/tools/filter - 按服务器筛选工具
 * <p>
 * 参考官方文档: <a href="https://docs.spring.io/spring-ai/reference/api/mcp/mcp-helpers.html">Spring AI MCP Utilities</a>
 */
@Slf4j
@CrossOrigin
@RestController
@RequestMapping("/api/v1/mcp")
@Tag(name = "MCP工具管理", description = "管理和查询MCP工具及服务器状态")
public class McpToolController {

    @Resource
    private List<McpSyncClient> mcpSyncClients;

    /**
     * 获取所有MCP服务器列表
     */
    @Operation(
            summary = "获取MCP服务器列表",
            description = "获取所有已注册的MCP服务器基本信息"
    )
    @GetMapping("/servers")
    public Map<String, Object> getServers() {
        log.info("查询MCP服务器列表");

        if (CollectionUtils.isEmpty(mcpSyncClients)) {
            return Map.of(
                    "success", false,
                    "message", "未找到已注册的MCP客户端",
                    "servers", Collections.emptyList()
            );
        }

        List<Map<String, Object>> servers = mcpSyncClients.stream()
                .map(client -> {
                    McpSchema.Implementation info = client.getServerInfo();
                    int toolCount;
                    try {
                        toolCount = client.listTools().tools().size();
                    } catch (Exception e) {
                        log.error("获取服务器[{}]的工具数量失败: {}", info.name(), e.getMessage());
                        toolCount = -1;
                    }
                    Map<String, Object> serverData = new HashMap<>();
                    serverData.put("name", info.name());
                    serverData.put("version", info.version());
                    serverData.put("toolCount", toolCount);
                    return serverData;
                })
                .toList();

        log.info("查询到 {} 个MCP服务器", servers.size());

        return Map.of(
                "success", true,
                "totalServers", servers.size(),
                "servers", servers
        );
    }

    /**
     * 获取所有MCP工具状态
     * 返回所有已注册的MCP服务器提供的工具列表
     */
    @Operation(
            summary = "获取所有MCP工具",
            description = "查询所有已注册的MCP工具列表及其详细信息，按服务器分组"
    )
    @GetMapping("/tools")
    public Map<String, Object> getAllTools() {
        log.info("查询所有MCP工具状态");

        if (CollectionUtils.isEmpty(mcpSyncClients)) {
            log.warn("未找到已注册的MCP客户端");
            return Map.of(
                    "success", false,
                    "message", "未找到已注册的MCP客户端",
                    "servers", Collections.emptyList()
            );
        }

        List<Map<String, Object>> serverInfoList = new ArrayList<>();
        int totalTools = 0;

        for (McpSyncClient client : mcpSyncClients) {
            try {
                // 获取服务器信息
                McpSchema.Implementation serverInfo = client.getServerInfo();
                String serverName = serverInfo.name();
                String serverVersion = serverInfo.version();

                // 获取工具列表
                List<McpSchema.Tool> tools = client.listTools().tools();
                totalTools += tools.size();

                Map<String, Object> serverData = new HashMap<>();
                serverData.put("serverName", serverName);
                serverData.put("serverVersion", serverVersion);
                serverData.put("toolCount", tools.size());
                serverData.put("tools", tools.stream().map(tool -> Map.of(
                        "name", tool.name(),
                        "description", tool.description() != null ? tool.description() : "",
                        "inputSchema", tool.inputSchema()
                )).collect(Collectors.toList()));

                serverInfoList.add(serverData);
                log.debug("从MCP服务器 [{}] 获取到 {} 个工具", serverName, tools.size());

            } catch (Exception e) {
                log.error("获取MCP工具列表失败: {}", e.getMessage(), e);
            }
        }

        log.info("共查询到 {} 个MCP服务器，总计 {} 个工具", serverInfoList.size(), totalTools);

        return Map.of(
                "success", true,
                "totalServers", serverInfoList.size(),
                "totalTools", totalTools,
                "servers", serverInfoList
        );
    }

    /**
     * 根据服务器名称筛选工具
     * 参考官方文档的动态选择功能
     */
    @Operation(
            summary = "按服务器名称筛选工具",
            description = "根据指定的服务器名称列表筛选并返回相应的MCP工具"
    )
    @GetMapping("/tools/filter")
    public Map<String, Object> getToolsByServers(
            @Parameter(description = "服务器名称列表，多个用逗号分隔，例如：filesystem,fetch")
            @RequestParam(value = "servers") String serverNames) {

        Set<String> allowedServerNames = Arrays.stream(serverNames.split(","))
                .map(String::trim)
                .collect(Collectors.toSet());

        log.info("按服务器名称筛选工具: {}", allowedServerNames);

        if (CollectionUtils.isEmpty(mcpSyncClients)) {
            return Map.of(
                    "success", false,
                    "message", "未找到已注册的MCP客户端"
            );
        }

        // 参考官方文档：动态选择MCP客户端子集
        List<McpSyncClient> selectedClients = mcpSyncClients.stream()
                .filter(c -> allowedServerNames.contains(c.getServerInfo().name()))
                .toList();

        if (selectedClients.isEmpty()) {
            log.warn("未找到匹配的MCP服务器: {}", allowedServerNames);
            return Map.of(
                    "success", false,
                    "message", "未找到匹配的服务器: " + allowedServerNames,
                    "availableServers", mcpSyncClients.stream()
                            .map(c -> c.getServerInfo().name())
                            .collect(Collectors.toList())
            );
        }

        List<Map<String, Object>> results = new ArrayList<>();
        int totalTools = 0;

        for (McpSyncClient client : selectedClients) {
            try {
                String serverName = client.getServerInfo().name();
                List<McpSchema.Tool> tools = client.listTools().tools();
                totalTools += tools.size();

                results.add(Map.of(
                        "serverName", serverName,
                        "toolCount", tools.size(),
                        "tools", tools.stream().map(tool -> Map.of(
                                "name", tool.name(),
                                "description", tool.description() != null ? tool.description() : ""
                        )).collect(Collectors.toList())
                ));

            } catch (Exception e) {
                log.error("获取工具列表失败: {}", e.getMessage(), e);
            }
        }

        log.info("筛选到 {} 个服务器，共 {} 个工具", results.size(), totalTools);

        return Map.of(
                "success", true,
                "selectedServers", allowedServerNames,
                "matchedCount", results.size(),
                "totalTools", totalTools,
                "results", results
        );
    }

    /**
     * 获取指定工具的详细信息
     */
    @Operation(
            summary = "获取工具详情",
            description = "根据工具名称获取详细的工具定义信息"
    )
    @GetMapping("/tools/{toolName}")
    public Map<String, Object> getToolDetail(
            @Parameter(description = "工具名称，例如：read_file")
            @PathVariable String toolName) {

        log.info("查询工具详情: {}", toolName);

        if (CollectionUtils.isEmpty(mcpSyncClients)) {
            return Map.of(
                    "success", false,
                    "message", "未找到已注册的MCP客户端"
            );
        }

        for (McpSyncClient client : mcpSyncClients) {
            try {
                List<McpSchema.Tool> tools = client.listTools().tools();
                Optional<McpSchema.Tool> tool = tools.stream()
                        .filter(t -> t.name().equals(toolName))
                        .findFirst();

                if (tool.isPresent()) {
                    McpSchema.Tool foundTool = tool.get();
                    log.info("找到工具: {}", toolName);
                    return Map.of(
                            "success", true,
                            "serverName", client.getServerInfo().name(),
                            "tool", Map.of(
                                    "name", foundTool.name(),
                                    "description", foundTool.description() != null ? foundTool.description() : "",
                                    "inputSchema", foundTool.inputSchema()
                            )
                    );
                }
            } catch (Exception e) {
                log.error("查询工具详情失败: {}", e.getMessage(), e);
            }
        }

        log.warn("未找到工具: {}", toolName);
        return Map.of(
                "success", false,
                "message", "未找到工具: " + toolName
        );
    }
}

